package tools.gff;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;


public class gff3Utils {

	final static String sep="\t";

	public static void main(String[] args) throws Exception{
		if(args.length>0){
			if(args[0].equals("toKnownGenes")&&args.length>1){
				toKnownGene(args[1]);
			}else if(args[0].equals("gmToStrainBED")&&args.length>3){
//				gmToStrainBED(args[1], args[2], args[3]);
			}else if(args[0].equals("overlapLines")&&args.length>2){
				overlapLines(args[1],args[2]);
			}else if(args[0].equals("selfOverlapClusters")&&args.length>1){
				selfOverlapClusters(args[1]);
			}else if(args[0].equals("removeSinglets")&&args.length>2){
				removeSinglets(args[1],args[2],args[3]);
			}else if(args[0].equals("printLines")&&args.length>2){
				printLines(args[1],args[2],true);
			}else{
				System.err.println(printHelp());
				System.exit(616);
			}
		}else{
			System.err.println(printHelp());
			System.exit(616);
		}
	}
	
	private static String printHelp(){
		String help="Usage: gff3Utils <cmd> <input>\n";
		help+="where <cmd> is:\n";
		help+="toKnownGenes - Converts a GFF3 file with genes to the known genes table from UCSC\n";
		help+="\t<input> = <GFF3 file>\n";
		help+="overlapLines - Check for overlap between the lines of two GFF3 files. Prints a file with line numbers\n";
		help+="\t<input> = <GFF3 file1> <GFF3 file2>\n";
		help+="selfOverlapClusters - Takes an overlap file from overlap lines (made by searching a file against itself) and generates a list of clusters\n";
		help+="\t<input> = <GFF3 self overlap file>\n";
		help+="removeSinglets - removes the singlets and prints the lines from the GFF file. Prints the singlet clusters to a separate file\n";
		help+="\t<input> = <GFF3 self cluster file> <GFF3 file> <outPrefix>\n";
		help+="printLines - prints the lines given by column 1 in the subfile\n";
		help+="\t<input> = <GFF3 file> <subFile>\n";
		
		
		return help;
	}
	
	public static void printLines(String gffFile,String subFile,boolean inSubFile)throws Exception{
		BufferedReader in= new BufferedReader(new FileReader(subFile));
		HashSet<Integer> lines= new HashSet<Integer>();
		String[] l;
		for(String s=in.readLine();s!=null;s=in.readLine()){
			l=s.split(sep);
			lines.add(Integer.parseInt(l[0]));
		}
		in.close();
		in= new BufferedReader(new FileReader(gffFile));
		int linenr=0;
		for(String s=in.readLine();s!=null;s=in.readLine()){
			if(lines.contains(linenr)==inSubFile){
				System.out.println(s);
			}
			++linenr;
		}
		in.close();
	}
	
	public static void removeSinglets(String selfClusterFile,String gffFile, String outPrefix)throws Exception{
		BufferedReader in= new BufferedReader(new FileReader(selfClusterFile));
		HashMap<String, HashSet<String>> clusters= new HashMap<String, HashSet<String>>();
		HashSet<String> clusterSet;
		String[] l;
		for(String s=in.readLine();s!=null;s=in.readLine()){
			l=s.split(sep);
			if(!clusters.containsKey(l[1])){
				clusters.put(l[1], new HashSet<String>());
			}
			clusters.get(l[1]).add(l[0]);
		}
		BufferedWriter out= new BufferedWriter(new FileWriter(outPrefix+"_singlets.csv"));
		HashMap<Integer, String> linesToPrint= new HashMap<Integer, String>();
		for (String cluster : clusters.keySet()) {
			clusterSet=clusters.get(cluster);
			if(clusterSet.size()==1){
				//push to singlet file
				for (String id : clusterSet) {
					out.write(id+sep+cluster+"\n");
				}
			}else{
				//prep for extraction from the GFF3 file
				for (String id : clusterSet) {
					linesToPrint.put(Integer.parseInt(id), cluster);
				}
			}
		}
		out.close();
		out= new BufferedWriter(new FileWriter(outPrefix+"_doublets.gffExt"));
		in= new BufferedReader(new FileReader(gffFile));
		int linenr=0;
		for(String s=in.readLine();s!=null;s=in.readLine()){
			if(linesToPrint.containsKey(linenr)){
				out.write(s+sep+linenr+sep+linesToPrint.get(linenr)+"\n");
			}
			++linenr;
		}
		out.close();
	}
	
	public static void selfOverlapClusters(String gffSelfOverlapfile)throws Exception{
		BufferedReader in= new BufferedReader(new FileReader(gffSelfOverlapfile));
		HashMap<String, Integer> idToCluster= new HashMap<String, Integer>();
		HashMap<Integer, HashSet<String>> clusters= new HashMap<Integer, HashSet<String>>();
		HashSet<String> tmp;
		String[] l;
		int cluster1,cluster2,nr=0;
		for(String s=in.readLine();s!=null;s=in.readLine()){
			l=s.split(sep);
			if(idToCluster.containsKey(l[0])){
				cluster1=idToCluster.get(l[0]);
				if(idToCluster.containsKey(l[1])){
					cluster2=idToCluster.get(l[1]);
					if(cluster1!=cluster2){
						//add cluster2 to cluster1
						for (String id : clusters.get(cluster2)) {
							idToCluster.put(id, cluster1);
						}
						clusters.get(cluster1).addAll(clusters.get(cluster2));
					}
				}else{
					//add l[1] to cluster1
					clusters.get(cluster1).add(l[1]);
					idToCluster.put(l[1], cluster1);
				}
			}else if(idToCluster.containsKey(l[1])){
				cluster2=idToCluster.get(l[1]);
				//add l[0] to cluster2
				clusters.get(cluster2).add(l[0]);
				idToCluster.put(l[0], cluster2);
			}else{
				//create new cluster
				tmp= new HashSet<String>();
				tmp.add(l[0]);
				idToCluster.put(l[0], nr);
				tmp.add(l[1]);
				idToCluster.put(l[1], nr);
				clusters.put(nr, tmp);
				++nr;
			}
		}
		//print
		for (String id : idToCluster.keySet()) {
			System.out.println(id+sep+idToCluster.get(id));
		}
	}
	
	public static void overlapLines(String gff3File1, String gff3File2)throws Exception{
		BufferedReader in=new BufferedReader(new FileReader(gff3File1));
		HashMap<String, ArrayList<gffLineInt>> file1= new HashMap<String, ArrayList<gffLineInt>>();
		gffLine gl;
		int maxLength=0;
		int linenr=0;
		for(String s=in.readLine();s!=null;s=in.readLine()){
			gl= new gffLine(s);
			if(!file1.containsKey(gl.getChr())){
				file1.put(gl.getChr(), new ArrayList<gffLineInt>());
			}
			maxLength=maxLength>gl.getEnd()-gl.getStart()?maxLength:gl.getEnd()-gl.getStart();
			
			file1.get(gl.getChr()).add(new gffLineInt(linenr, gl));
			linenr++;
		}
		//Sorting
		for (ArrayList<gffLineInt> list : file1.values()) {
			Collections.sort(list, new gffLineInt_comparator());
		}
		in.close();
		in=new BufferedReader(new FileReader(gff3File2));
		linenr=0;
		for(String s=in.readLine();s!=null;s=in.readLine()){
			gl= new gffLine(s);
			if(file1.containsKey(gl.getChr())){
				for (gffLineInt gl1 : file1.get(gl.getChr())) {
					if(gl1.overlap(gl)){
						System.out.println(gl1.getLine()+sep+linenr);
					}
				}
			}
			linenr++;
		}
	}
	
	public static void toKnownGene(String gff3File)throws Exception{
		ArrayList<GFF3node> genes=new ArrayList<GFF3node>();
		ArrayList<gffLine> exons;
		HashMap<String, GFF3node> parents= new HashMap<String, GFF3node>();
		
		BufferedReader in= new BufferedReader(new FileReader(gff3File));
		gffLine curLine;
		GFF3node curNode,tmpNode;
		HashSet<String> toHandle= new HashSet<String>();
		toHandle.add("gene");
		toHandle.add("transposable_element_gene");
		toHandle.add("mRNA");
		toHandle.add("miRNA");
		toHandle.add("snRNA");
		toHandle.add("rRNA");
		toHandle.add("snoRNA");
		toHandle.add("ncRNA");
		toHandle.add("tRNA");
		toHandle.add("exon");
		toHandle.add("protein");
		HashSet<String> isGene= new HashSet<String>();
		isGene.add("gene");
		isGene.add("transposable_element_gene");
		HashSet<String> toPrint= new HashSet<String>();
		toPrint.add("mRNA");
		toPrint.add("miRNA");
		toPrint.add("snRNA");
		toPrint.add("rRNA");
		toPrint.add("snoRNA");
		toPrint.add("ncRNA");
		toPrint.add("tRNA");
		Iterator<String> p;
		Iterator<GFF3node> iterateGFF3nodes,iterateChildren;
		String line,starts,ends;
		final Comparator<gffLine> comp = new gffLine_startComparator();
		for(String s=in.readLine();s!=null;s=in.readLine()){
			curLine=new gffLine(s);
			if(toHandle.contains(curLine.getFeature())){
				curNode= new GFF3node(curLine);
				if(isGene.contains(curLine.getFeature())){
					genes.add(curNode);
				}
				if(curLine.hasID()){
					parents.put(curLine.getID(), curNode);
				}
				if(curLine.hasParent()){
					p=curLine.iterateParents();
					for (;p.hasNext();) {
						final String parent = p.next();
						if(parents.containsKey(parent)){
							parents.get(parent).addChild(curNode);
						}else{
							//Assumes parents are defined before the derivatives or children
							System.err.println("Parent "+parent+" has to be defined before called.");
						}
					}
				}
				if(curLine.hasDerives_from()){
					p= curLine.iterateDerives_from();
					for(;p.hasNext();){
						final String source= p.next();
						if(parents.containsKey(source)){
							parents.get(source).addDerivative(curNode);
						}else{
							//Assumes parents are defined before the derivatives or children
//							System.err.println("The source "+source+" has to be defined before called.");
						}
					}
				}
			}
		}
		//print
		//header
		System.out.println("#name"+sep+"chrom"+sep+"strand"+sep+"txStart"+sep+"txEnd"+sep+"cdsStart"+sep+"cdsEnd"+sep+"exonCount"+sep+"exonStarts"+sep+"exonEnds"+sep+"proteinID"+sep+"alignID");
		//body
		for (GFF3node gene : genes) {
			iterateGFF3nodes=gene.iterateChildren();
			for(;iterateGFF3nodes.hasNext();){
				curNode=iterateGFF3nodes.next();
				curLine=curNode.getCur();
				if(toPrint.contains(curLine.getFeature())){
					//maybe I should use the gene ID in the first column?
					line=curLine.getID()+sep+curLine.getChr()+sep+curLine.getStrand()+sep+(curLine.getStart()-1)+sep+curLine.getEnd();
					if(curNode.derivativesSize()==0){
						//When no protein is derived from the mRNA, the start and end is used as CDS values
						line+=sep+(curLine.getStart()-1)+sep+curLine.getEnd();
					}else if(curNode.derivativesSize()==1){
						line+=sep+(curNode.getFirstDerivative().getCur().getStart()-1)+sep+curNode.getFirstDerivative().getCur().getEnd();
					}else{
						System.err.println("multiple proteins associated with mRNA, "+curLine.getID()+", only using first.");
						line+=sep+(curNode.getFirstDerivative().getCur().getStart()-1)+sep+curNode.getFirstDerivative().getCur().getEnd();
					}
					exons=new ArrayList<gffLine>();
					iterateChildren=curNode.iterateChildren();
					for(;iterateChildren.hasNext();){
						tmpNode=iterateChildren.next();
						if(tmpNode.getCur().getFeature().equals("exon")){
							exons.add(tmpNode.getCur());
						}
					}
					Collections.sort(exons, comp);
					line+=sep+exons.size();
					starts="";
					ends="";
					for (gffLine exon : exons) {
						starts+=(exon.getStart()-1)+",";
						ends+=exon.getEnd()+",";
					}
					line+=sep+starts+sep+ends+sep+""+sep+curLine.getID();
					System.out.println(line);
				}else{
					System.err.println("gene "+gene.getCur().getID()+" has a child of type: "+curLine.getFeature());
				}
			}
		}
	}
	
//	public static String toKnownGene_toString(GFF3node curNode){
//		final
//		String line=
//	}
}

class gffLineInt{
	private int line;
	private gffLine gl;
	
	protected gffLineInt(int line, gffLine gl) {
		super();
		this.line = line;
		this.gl = gl;
	}

	protected int getLine() {
		return line;
	}

	protected gffLine getGl() {
		return gl;
	}
	
	protected boolean overlap(gffLine gl2){
		return this.getGl().overlap(gl2);
	}
}

class gffLineInt_comparator implements Comparator<gffLineInt>{

	public int compare(gffLineInt o1, gffLineInt o2) {
		// TODO Auto-generated method stub
		return o1.getGl().getStart()-o2.getGl().getStart();
	}
	
}